import torch
import torch.onnx
import sys
import models.barcode_model_small as models


# Load the PyTorch model
model_path = r"C:\Users\Administrator\Desktop\barcode_train\runs_barcode_1d\runs2\sigma_2epoch_30.pt"
device = torch.device("cpu")
model = models.GrayDemoNet2()

checkpoint = torch.load(model_path, map_location=device)
if isinstance(checkpoint, dict):
    model.load_state_dict(checkpoint["model_state_dict"])
else:
    model = checkpoint

model.eval()
model.to(device)
onnx_model_path = model_path.replace(".pt", ".onnx")
input = torch.randn(1, 1, 128, 128).to(device)

# Export the model to ONNX
# The "model.onnx" here is the name of the exported ONNX file, you can modify it as needed
# input_names and output_names can be set as needed to identify the names of the model's input and output layers
torch.onnx.export(
    model,
    input,
    onnx_model_path,
    export_params=True,  # 导出模型中的学习到的参数
    opset_version=11,
    do_constant_folding=True,  # 是否执行常量折叠优化
    input_names=["input"],
    output_names=["output"],
)
